# 创建数字数据集示例
import cv2
import numpy as np
import torch
import os
from torch.utils.data import Dataset

class DigitDataset(Dataset):
    def __init__(self, root_dir):
        self.samples = []
        for label in os.listdir(root_dir):
            for img_file in os.listdir(f"{root_dir}/{label}"):
                path = f"{root_dir}/{label}/{img_file}"
                self.samples.append((path, label))
                
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        path, label = self.samples[idx]
        img = cv2.imread(path, 0)  # 灰度图
        img = cv2.resize(img, (32, 32))
        img = img.astype(np.float32) / 255.0
        img = torch.FloatTensor(img).unsqueeze(0)  # 添加通道维度
        
        # 转换标签为数字索引（需自定义字符集）
        char_set = "0123456789."
        target = [char_set.index(c) for c in label]
        return img, torch.LongTensor(target)
